from typing import Dict
import pymysql
from db_utils import con_my_sql
from db_decorator_key_apply import private_key_crypto_decorator
from config import config_instance


# 创建密钥管理数据库连接
key_conn = pymysql.connect(**config_instance.key_conn_config)

@private_key_crypto_decorator
def private_con_my_sql(sql_code, params=None, connection=None):
    return con_my_sql(connection, sql_code, params)


def insert_sm2_key_pairs(key_pair: Dict):
    """
    将传入的包含 SM2 公私钥对信息的字典插入到 sm2_key_pairs 表中。
    注意：这里需要调用者确保传入的字典中包含合法的 'id' 字段值，且该值在表中具有唯一性。

    :param key_pair: 包含 SM2 公私键对信息以及 'id' 字段的字典
    :return: 如果插入成功返回 True，出现异常返回 False
    """
    try:
        query = "INSERT INTO sm2_key_pairs (id, encrypted_signature_private_key, signature_public_key, encrypted_encryption_private_key, encryption_public_key, signature_private_key_iv, encryption_private_key_iv) VALUES (%(id)s, %(encrypted_signature_private_key)s, %(signature_public_key)s, %(encrypted_encryption_private_key)s, %(encryption_public_key)s, %(signature_private_key_iv)s, %(encryption_private_key_iv)s)"
        params = key_pair.copy()  # 直接使用传入字典的副本作为参数，包含 'id' 及其他键值对
        private_con_my_sql(query, params, key_conn)
        return True
    except Exception as e:
        print("插入 SM2 公私键对数据库类错误：", str(e))
        return False


def update_sm2_key_pairs(key_pair_id: int, updated_key_pair: Dict):
    """
    根据传入的键对 ID 和更新后的键对信息字典更新 sm2_key_pairs 表中的记录。

    :param key_pair_id: 要更新的键对 ID
    :param updated_key_pair: 包含更新后 SM2 公私键对信息的字典
    :return: 如果更新成功返回 True，出现异常返回 False
    """
    try:
        update_fields = []
        params = {}
        for key, value in updated_key_pair.items():
            if key in ['encrypted_signature_private_key', 'signature_public_key', 'encrypted_encryption_private_key', 'encryption_public_key','signature_private_key_iv', 'encryption_private_key_iv']:
                update_fields.append(f"{key} = %({key})s")
                params[key] = value
        query = f"UPDATE sm2_key_pairs SET {','.join(update_fields)} WHERE id = {key_pair_id}"
        private_con_my_sql(query, params, key_conn)
        return True
    except Exception as e:
        print("更新 SM2 公私键对数据库类错误：", str(e))
        return False


def get_sm2_key_pairs(key_pair_id: int = None):
    """
    根据传入的键对 ID 获取 sm2_key_pairs 表中的记录，如果未传入 ID，则获取所有记录。

    :param key_pair_id: 可选的键对 ID，默认为 None
    :return: 返回查询到的键对信息字典列表，如果出现异常返回 None
    """
    try:
        if key_pair_id:
            query = "SELECT * FROM sm2_key_pairs WHERE id = %s"
            result = private_con_my_sql(query, (key_pair_id,), key_conn)
        else:
            query = "SELECT * FROM sm2_key_pairs"
            result = private_con_my_sql(query, None, key_conn)
        return result
    except Exception as e:
        print("获取 SM2 公私键对数据库类错误：", str(e))
        return None


# 测试数据
# 签名
# 私钥：45ebc875cac44a2420eeace261423df663a1f0378c68a26c932e409139d73451
# 公钥：95371cf1c49dc4127bc419622cce2023cfb19227c20517ffc0378f63d0670dd398b97225b18e7f7fc1233bd2eb84c5e0e04b4b1c99ae06bbe4c89b18ce858c26
# 加密
# 私钥：5d1c38bc16b7a69ceb9cfa9cd7026e7011e5568f33f5223f36c767543ac7224c
# 公钥：d6633100664750842ea5ac89f329525d955130b66a0297c57b0f3ea53a44427f64b104796418f480499e6e0f4efdc95b25b9358c3d99f5ba67b06c0b0b431894


def apply_test_data():
    key_pair_id = 1
    existing_key_pairs = get_sm2_key_pairs(key_pair_id)
    if existing_key_pairs:
        # 如果存在id为1的记录，执行更新操作
        signature_key_pair = {
            "encrypted_signature_private_key": bytes.fromhex(
                "45ebc875cac44a2420eeace261423df663a1f0378c68a26c932e409139d73451"),
            "signature_public_key": bytes.fromhex(
                "95371cf1c49dc4127bc419622cce2023cfb19227c20517ffc0378f63d0670dd398b97225b18e7f7fc1233bd2eb84c5e0e04b4b1c99ae06bbe4c89b18ce858c26"),
            "encrypted_encryption_private_key": bytes.fromhex(
                "5d1c38bc16b7a69ceb9cfa9cd7026e7011e5568f33f5223f36c767543ac7224c"),
            "encryption_public_key": bytes.fromhex(
                "d6633100664750842ea5ac89f329525d955130b66a0297c57b0f3ea53a44427f64b104796418f480499e6e0f4efdc95b25b9358c3d99f5ba67b06c0b0b431894"),
            "signature_private_key_iv": b'11',
            "encryption_private_key_iv": b'11'
        }
        update_result = update_sm2_key_pairs(key_pair_id, signature_key_pair)
        if update_result:
            print("公私钥对更新操作成功")
        else:
            print("公私钥对更新操作失败")
    else:
        # 如果不存在id为1的记录，执行插入操作
        signature_key_pair = {
            "id": 1,
            "encrypted_signature_private_key": bytes.fromhex(
                "45ebc875cac44a2420eeace261423df663a1f0378c68a26c932e409139d73451"),
            "signature_public_key": bytes.fromhex(
                "95371cf1c49dc4127bc419622cce2023cfb19227c20517ffc0378f63d0670dd398b97225b18e7f7fc1233bd2eb84c5e0e04b4b1c99ae06bbe4c89b18ce858c26"),
            "encrypted_encryption_private_key": bytes.fromhex(
                "5d1c38bc16b7a69ceb9cfa9cd7026e7011e5568f33f5223f36c767543ac7224c"),
            "encryption_public_key": bytes.fromhex(
                "d6633100664750842ea5ac89f329525d955130b66a0297c57b0f3ea53a44427f64b104796418f480499e6e0f4efdc95b25b9358c3d99f5ba67b06c0b0b431894"),
            "signature_private_key_iv": b'11',
            "encryption_private_key_iv": b'11'
        }
        insert_signature_result = insert_sm2_key_pairs(signature_key_pair)
        if insert_signature_result:
            print("公私钥对插入操作成功")
        else:
            print("公私钥对插入操作失败")

    # 获取所有记录，验证操作是否成功（简单查看）
    all_key_pairs = get_sm2_key_pairs()
    if all_key_pairs:
        print("获取所有记录成功，记录信息:", all_key_pairs)
    else:
        print("获取所有记录失败，无法验证操作情况")

def main():
    # 测试插入操作
    test_key_pair = {
        "id": 1,  # 这里手动指定一个合法的 'id' 值，实际应用中需确保唯一性
        "encrypted_signature_private_key": b'\x01\x02\x03',  # 示例二进制数据，实际应替换为真实的加密私钥数据
        "signature_public_key": b'\x04\x05\x06',
        "encrypted_encryption_private_key": b'\x07\x08\x09',
        "encryption_public_key": b'\x10\x11\x12',
        "signature_private_key_iv": b'\x13\x14\x15',
        "encryption_private_key_iv": b'\x16\x17\x18'
    }
    insert_result = insert_sm2_key_pairs(test_key_pair)
    if insert_result:
        print("插入操作成功")
    else:
        print("插入操作失败")

    # 获取插入后的所有记录，验证插入是否成功（假设插入成功后获取最新插入的那条记录进行后续更新测试，这里简单获取所有记录查看）
    all_key_pairs = get_sm2_key_pairs()
    if all_key_pairs:
        target_key_pair_id = 1  # 这里指定之前插入时使用的 'id' 值，用于更新测试
        print("获取所有记录成功，准备用于更新测试的记录ID:", target_key_pair_id)

        # 测试更新操作
        updated_key_pair = {
            "encrypted_signature_private_key": b'\x19\x20\x21',  # 更新后的示例二进制数据
            "signature_public_key": b'\x22\x23\x24'
        }
        update_result = update_sm2_key_pairs(target_key_pair_id, updated_key_pair)
        if update_result:
            print("更新操作成功")
        else:
            print("更新操作失败")

        # 再次获取记录，验证更新是否生效
        updated_key_pairs = get_sm2_key_pairs(target_key_pair_id)
        if updated_key_pairs:
            print("更新后获取记录成功，记录信息:", updated_key_pairs[0])
        else:
            print("更新后获取记录失败")
    else:
        print("获取所有记录失败，无法进行更新测试")

